package demo.webauthn.controller;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.yubico.webauthn.*;
import com.yubico.webauthn.data.*;
import com.yubico.webauthn.exception.AssertionFailedException;
import com.yubico.webauthn.exception.RegistrationFailedException;
import demo.webauthn.domain.Result;
import demo.webauthn.domain.WebAuthUser;
import demo.webauthn.mapper.WebAuthUserMapper;
import demo.webauthn.util.RedisUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.*;

import javax.servlet.http.HttpSession;
import java.io.IOException;
import java.security.SecureRandom;
import java.time.Duration;
import java.util.List;
import java.util.Map;

/**
 * WebAuthn 控制器
 */
@RestController
@RequestMapping("/webauthn")
public class WebAuthnDemoController {

    @Autowired
    private RelyingParty relyingParty;
    @Autowired
    private WebAuthUserMapper webAuthUserMapper;

    @GetMapping("/test")
    public Result<String> test() {
        return Result.ok("测试成功");
    }

    /**
     * 用户注册
     */
    @PostMapping("/register")
    public Result<String> register(@RequestBody Map<String, String> map, HttpSession session) throws JsonProcessingException {

        String userName = map.get("userName");
        String displayName = map.get("displayName");
        if (StringUtils.isEmpty(userName) || StringUtils.isEmpty(displayName)) {
            return Result.fail("用户名不能为空");
        }

        List<WebAuthUser> webAuthUsers = webAuthUserMapper.selectByUsername(userName);
        if (!CollectionUtils.isEmpty(webAuthUsers)) {
            return Result.fail("用户已被注册");
        }

        byte[] bytes = new byte[32];
        new SecureRandom().nextBytes(bytes);
        UserIdentity userIdentity = UserIdentity.builder()
                .name(userName)
                .displayName(displayName)
                .id(new ByteArray(bytes))
                .build();
        StartRegistrationOptions registrationOptions = StartRegistrationOptions.builder()
                .user(userIdentity)
                .build();
        PublicKeyCredentialCreationOptions registration = relyingParty.startRegistration(registrationOptions);
        RedisUtil.setCacheObject("register-" + userName, registration.toJson(), Duration.ofHours(48));

        try {
            String credentialsCreateJson = registration.toCredentialsCreateJson();
            return Result.ok("注册成功", credentialsCreateJson);
        } catch (JsonProcessingException e) {
            e.printStackTrace();
            return Result.fail("注册失败");
        }
    }

    /**
     * 注册认证
     */
    @PostMapping("/registerAuth")
    public Result<String> registerAuth(@RequestBody Map<String, String> map, HttpSession session) {
        String username = map.get("userName");
        String credential = map.get("credential");
        try {
            Object o = RedisUtil.getCacheObject("register-" + username);
            if (o == null) return Result.fail("缓存请求已过期, 请尝试再次注册");

            PublicKeyCredentialCreationOptions requestOptions = PublicKeyCredentialCreationOptions.fromJson(o.toString());
            if (requestOptions != null) {
                PublicKeyCredential<AuthenticatorAttestationResponse, ClientRegistrationExtensionOutputs> pkc =
                        PublicKeyCredential.parseRegistrationResponseJson(credential);

                FinishRegistrationOptions options = FinishRegistrationOptions.builder()
                        .request(requestOptions)
                        .response(pkc)
                        .build();

                RegistrationResult result = relyingParty.finishRegistration(options);
                String credentialId = result.getKeyId().getId().getBase64();
                String publicKey = result.getPublicKeyCose().getBase64();
                WebAuthUser webauthUser = new WebAuthUser();
                webauthUser.setCredentialId(credentialId);
                webauthUser.setDisplayName(username);
                webauthUser.setPublicKey(publicKey);
                webauthUser.setAlgorithm(requestOptions.getUser().getId().getBase64());
                webauthUser.setUserName(username);
                webAuthUserMapper.insert(webauthUser);

                return Result.ok("操作成功");
            } else {
                return Result.fail("缓存请求已过期, 请尝试再次注册");
            }
        } catch (RegistrationFailedException | IOException e) {
            e.printStackTrace();
            return Result.fail("操作失败");
        }

    }

    /**
     * 用户登录
     */
    @PostMapping("/login")
    public Result<String> login(@RequestBody Map<String, String> map, HttpSession session) {
        String username = map.get("userName");
        List<WebAuthUser> webAuthUsers = webAuthUserMapper.selectByUsername(username);
        if (CollectionUtils.isEmpty(webAuthUsers)) {
            Result.fail("用户名不存在");
        }
        AssertionRequest request = relyingParty.startAssertion(StartAssertionOptions.builder()
                .username(username)
                .build());
        try {
            RedisUtil.setCacheObject("login-" + username, request.toJson());
            return Result.ok("操作成功", request.toCredentialsGetJson());
        } catch (JsonProcessingException e) {
            e.printStackTrace();
            return Result.fail("操作失败");
        }
    }

    /**
     * 完成登录
     */
    @PostMapping("/finishLogin")
    public Result<String> finishLogin(@RequestBody Map<String, String> map, HttpSession session) {

        String username = map.get("userName");
        String credential = map.get("credential");

        try {
            PublicKeyCredential<AuthenticatorAssertionResponse, ClientAssertionExtensionOutputs> pkc;
            pkc = PublicKeyCredential.parseAssertionResponseJson(credential);
            Object o = RedisUtil.getCacheObject("login-" + username);
            AssertionRequest request = AssertionRequest.fromJson(o.toString());
            AssertionResult result = relyingParty.finishAssertion(FinishAssertionOptions.builder()
                    .request(request)
                    .response(pkc)
                    .build());
            if (result.isSuccess()) {
                return Result.ok("认证成功");
            } else {
                return Result.fail("认证失败, 请先注册");
            }
        } catch (IOException | AssertionFailedException | NullPointerException | IndexOutOfBoundsException e) {
            e.printStackTrace();
            return Result.fail("认证失败, 请先注册");
        }
    }

}
